package com.develophelper.aliyun.rocketmq.support;

import com.aliyun.openservices.ons.api.*;
import com.aliyun.openservices.ons.api.exception.ONSClientException;
import com.aliyun.openservices.ons.api.order.ConsumeOrderContext;
import com.aliyun.openservices.ons.api.order.MessageOrderListener;
import com.aliyun.openservices.ons.api.order.OrderAction;
import com.aliyun.openservices.ons.api.order.OrderConsumer;
import com.aliyun.openservices.shade.com.alibaba.rocketmq.client.exception.MQClientException;
import com.aliyun.openservices.shade.com.alibaba.rocketmq.common.protocol.heartbeat.MessageModel;
import com.develophelper.aliyun.rocketmq.annotation.ConsumeMode;
import com.develophelper.aliyun.rocketmq.annotation.RocketMQMessageListener;
import com.develophelper.aliyun.rocketmq.core.RocketMQListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.aop.framework.AopProxyUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanNameAware;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.SmartLifecycle;
import org.springframework.core.MethodParameter;
import org.springframework.messaging.converter.MessageConverter;
import org.springframework.messaging.converter.SmartMessageConverter;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.lang.reflect.Method;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.nio.charset.Charset;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author snaildev
 * @date 2021/05/13
 **/
@SuppressWarnings("WeakerAccess")
public class DefaultRocketMQListenerContainer implements InitializingBean,
        RocketMQListenerContainer, SmartLifecycle, ApplicationContextAware, BeanNameAware {
    private final static Logger log = LoggerFactory.getLogger(DefaultRocketMQListenerContainer.class);

    private ApplicationContext applicationContext;

    /**
     * The name of the DefaultRocketMQListenerContainer instance
     */
    private String name;

    private String nameServer;

    private String accessKey;

    private String secretKey;

    private Integer maxReconsumeTimes;

    private String consumerGroup;

    private String topic;

    private int consumeThreadMax = 64;

    private boolean running;

    // The following properties came from @RocketMQMessageListener.
    private ConsumeMode consumeMode;
    private ExpressionType selectorType;
    private String selectorExpression;
    private MessageModel messageModel;
    private long consumeTimeout;

    private String charset = "UTF-8";

    private MessageConverter messageConverter;

    private ConcurrentHashMap<String, List<RocketMQListenerExt>> listenerMap = new ConcurrentHashMap<>();

    private Set<String> traceIds = new LinkedHashSet<>();

    private Admin consumer;

    public String getNameServer() {
        return nameServer;
    }

    public void setNameServer(String nameServer) {
        this.nameServer = nameServer;
    }

    public void setConsumerGroup(String consumerGroup) {
        this.consumerGroup = consumerGroup;
    }

    public void setTopic(String topic) {
        this.topic = topic;
    }

    public MessageConverter getMessageConverter() {
        return messageConverter;
    }

    public DefaultRocketMQListenerContainer setMessageConverter(MessageConverter messageConverter) {
        this.messageConverter = messageConverter;
        return this;
    }

    public void addRocketMQListener(String selectorExpression, RocketMQListener rocketMQListener) {
        if (listenerMap.containsKey(selectorExpression)) {
            this.listenerMap.get(selectorExpression).add(new RocketMQListenerExt(rocketMQListener, getMessageType(rocketMQListener), getMethodParameter(rocketMQListener)));
        } else {
            this.listenerMap.put(selectorExpression, Arrays.asList(new RocketMQListenerExt(rocketMQListener, getMessageType(rocketMQListener), getMethodParameter(rocketMQListener))));
        }
    }

    public void setRocketMQMessageListener(RocketMQMessageListener anno) {
        this.consumeMode = anno.consumeMode();
        this.consumeThreadMax = anno.consumeThreadMax();
        this.messageModel = anno.messageModel();
        this.selectorType = anno.selectorType();
        this.selectorExpression = anno.selectorExpression();
        this.consumeTimeout = anno.consumeTimeout();
        this.accessKey = anno.accessKey();
        this.secretKey = anno.secretKey();
        this.maxReconsumeTimes = anno.maxReconsumeTimes();
        this.traceIds.add(anno.traceId());
    }

    @Override
    public void destroy() {
        this.setRunning(false);
        if (Objects.nonNull(consumer)) {
            consumer.shutdown();
        }
        log.info("container destroyed, {}", this.toString());
    }

    @Override
    public boolean isAutoStartup() {
        return true;
    }

    @Override
    public void stop(Runnable callback) {
        stop();
        callback.run();
    }

    @Override
    public void start() {
        if (this.isRunning()) {
            throw new IllegalStateException("container already running. " + this.toString());
        }

        try {
            consumer.start();
        } catch (ONSClientException e) {
            throw new IllegalStateException("Failed to start RocketMQ push consumer", e);
        }
        this.setRunning(true);

        log.info("running container: {}", this.toString());
    }

    @Override
    public void stop() {
        if (this.isRunning()) {
            if (Objects.nonNull(consumer)) {
                consumer.shutdown();
            }
            setRunning(false);
        }
    }

    @Override
    public boolean isRunning() {
        return running;
    }

    private void setRunning(boolean running) {
        this.running = running;
    }

    @Override
    public int getPhase() {
        // Returning Integer.MAX_VALUE only suggests that
        // we will be the first bean to shutdown and last bean to start
        return Integer.MAX_VALUE;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        this.selectorExpression = getSelectorExpression();
        initRocketMQPushConsumer();
    }

    private String getSelectorExpression() {
        Set<String> expressionSet = new HashSet<>();
        Enumeration<String> keys = this.listenerMap.keys();
        while (keys.hasMoreElements()) {
            String expression = keys.nextElement();
            for (String tag : Arrays.asList(expression.split("\\|\\|"))) {
                expressionSet.add(tag);
            }
        }
        return expressionSet.contains("*") ? "*" : String.join(" || ", expressionSet);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Override
    public String toString() {
        return "DefaultRocketMQListenerContainer{" +
                "consumerGroup='" + consumerGroup + '\'' +
                ", nameServer='" + nameServer + '\'' +
                ", topic='" + topic + '\'' +
                ", consumeMode=" + consumeMode +
                ", selectorType=" + selectorType +
                ", selectorExpression='" + selectorExpression + '\'' +
                ", messageModel=" + messageModel +
                '}';
    }

    public String getName() {
        return this.name;
    }

    @Override
    public void setBeanName(String name) {
        this.name = name;
    }

    public void setName(String name) {
        this.name = name;
    }

    private class RocketMQListenerExt {

        private RocketMQListener listener;

        private Type messageType;

        private MethodParameter methodParameter;

        public RocketMQListenerExt(RocketMQListener listener, Type messageType, MethodParameter methodParameter) {
            this.listener = listener;
            this.messageType = messageType;
            this.methodParameter = methodParameter;
        }

        public RocketMQListener getListener() {
            return listener;
        }

        public Type getMessageType() {
            return messageType;
        }

        public MethodParameter getMethodParameter() {
            return methodParameter;
        }
    }

    public class DefaultMessageListenerConcurrently implements MessageListener {
        @SuppressWarnings("unchecked")
        @Override
        public Action consume(Message message, ConsumeContext context) {
            try {
                traceIds.forEach(x -> MDC.put(x, Optional.ofNullable(message.getUserProperties(x)).orElse(UUID.randomUUID().toString())));
                log.debug("received msg: {}", message);
                long now = System.currentTimeMillis();
                handleMessage(message);
                long costTime = System.currentTimeMillis() - now;
                log.debug("consume {} cost: {} ms", message.getMsgID(), costTime);
            } catch (Exception e) {
                log.warn("consume message failed. messageId:{}, topic:{}, reconsumeTimes:{}", message.getMsgID(), message.getTopic(), message.getReconsumeTimes(), e);
                return Action.ReconsumeLater;
            } finally {
                traceIds.forEach(x -> MDC.remove(x));
            }

            return Action.CommitMessage;
        }
    }

    public class DefaultMessageListenerOrderly implements MessageOrderListener {
        @SuppressWarnings("unchecked")
        @Override
        public OrderAction consume(Message message, ConsumeOrderContext context) {
            try {
                traceIds.forEach(x -> MDC.put(x, Optional.ofNullable(message.getUserProperties(x)).orElse(UUID.randomUUID().toString())));
                log.debug("received msg: {}", message);
                long now = System.currentTimeMillis();
                handleMessage(message);
                long costTime = System.currentTimeMillis() - now;
                log.debug("consume {} cost: {} ms", message.getMsgID(), costTime);
            } catch (Exception e) {
                log.warn("consume message failed. messageId:{}, topic:{}, reconsumeTimes:{}", message.getMsgID(), message.getTopic(), message.getReconsumeTimes(), e);
                return OrderAction.Suspend;
            } finally {
                traceIds.forEach(x -> MDC.remove(x));
            }

            return OrderAction.Success;
        }
    }


    private void handleMessage(
            Message message) {
        if (!CollectionUtils.isEmpty(listenerMap)) {
            listenerMap.forEach((k, v) -> {
                if (k.equals("*") ||
                        Arrays.asList(k.split("\\|\\|")).contains(message.getTag())) {
                    for (RocketMQListenerExt rocketMQListenerExt : v) {
                        rocketMQListenerExt.getListener().onMessage(doConvertMessage(message, rocketMQListenerExt.getMessageType(), rocketMQListenerExt.getMethodParameter()));
                    }
                }
            });
        }
    }

    @SuppressWarnings("unchecked")
    private Object doConvertMessage(Message message, Type messageType, MethodParameter methodParameter) {
        if (Objects.equals(messageType, Message.class)) {
            return message;
        } else {
            String str = new String(message.getBody(), Charset.forName(charset));
            if (Objects.equals(messageType, String.class)) {
                return str;
            } else {
                // If msgType not string, use objectMapper change it.
                try {
                    if (messageType instanceof Class) {
                        //if the messageType has not Generic Parameter
                        return this.getMessageConverter().fromMessage(MessageBuilder.withPayload(str).build(), (Class<?>) messageType);
                    } else {
                        //if the messageType has Generic Parameter, then use SmartMessageConverter#fromMessage with third parameter "conversionHint".
                        //we have validate the MessageConverter is SmartMessageConverter in this#getMethodParameter.
                        return ((SmartMessageConverter) this.getMessageConverter()).fromMessage(MessageBuilder.withPayload(str).build(), (Class<?>) ((ParameterizedType) messageType).getRawType(), methodParameter);
                    }
                } catch (Exception e) {
                    log.info("convert failed. str:{}, msgType:{}", str, messageType);
                    throw new RuntimeException("cannot convert message to " + messageType, e);
                }
            }
        }
    }

    private MethodParameter getMethodParameter(RocketMQListener rocketMQListener) {
        Class<?> targetClass = AopProxyUtils.ultimateTargetClass(rocketMQListener);
        Type messageType = this.getMessageType(rocketMQListener);
        Class clazz = null;
        if (messageType instanceof ParameterizedType && messageConverter instanceof SmartMessageConverter) {
            clazz = (Class) ((ParameterizedType) messageType).getRawType();
        } else if (messageType instanceof Class) {
            clazz = (Class) messageType;
        } else {
            throw new RuntimeException("parameterType:" + messageType + " of onMessage method is not supported");
        }
        try {
            final Method method = targetClass.getMethod("onMessage", clazz);
            return new MethodParameter(method, 0);
        } catch (NoSuchMethodException e) {
            e.printStackTrace();
            throw new RuntimeException("parameterType:" + messageType + " of onMessage method is not supported");
        }
    }

    private Type getMessageType(RocketMQListener rocketMQListener) {
        Class<?> targetClass = AopProxyUtils.ultimateTargetClass(rocketMQListener);

        Type matchedGenericInterface = null;
        while (Objects.nonNull(targetClass)) {
            Type[] interfaces = targetClass.getGenericInterfaces();
            if (Objects.nonNull(interfaces)) {
                for (Type type : interfaces) {
                    if (type instanceof ParameterizedType &&
                            (Objects.equals(((ParameterizedType) type).getRawType(), RocketMQListener.class))) {
                        matchedGenericInterface = type;
                        break;
                    }
                }
            }
            targetClass = targetClass.getSuperclass();
        }
        if (Objects.isNull(matchedGenericInterface)) {
            return Object.class;
        }

        Type[] actualTypeArguments = ((ParameterizedType) matchedGenericInterface).getActualTypeArguments();
        if (Objects.nonNull(actualTypeArguments) && actualTypeArguments.length > 0) {
            return actualTypeArguments[0];
        }
        return Object.class;
    }

    private void initRocketMQPushConsumer() throws MQClientException {
        if (CollectionUtils.isEmpty(listenerMap)) {
            throw new IllegalArgumentException("Property 'listenerMap' is required");
        }
        Assert.notNull(consumerGroup, "Property 'consumerGroup' is required");
        Assert.notNull(nameServer, "Property 'nameServer' is required");
        Assert.notNull(topic, "Property 'topic' is required");

        Properties properties = getProperties();
        if (ConsumeMode.CONCURRENTLY.equals(consumeMode)) {
            consumer = ONSFactory.createConsumer(properties);
            switch (selectorType) {
                case TAG:
                    ((Consumer) consumer).subscribe(topic, selectorExpression, new DefaultMessageListenerConcurrently());
                    break;
//                case SQL92:
//                    ((Consumer) consumer).subscribe(topic, MessageSelector.bySql(selectorExpression), new DefaultMessageListenerConcurrently());
//                    break;
                default:
                    throw new IllegalArgumentException("Property 'selectorType' was wrong.");
            }
        } else if (ConsumeMode.ORDERLY.equals(consumeMode)) {
            consumer = ONSFactory.createOrderedConsumer(properties);
            switch (selectorType) {
                case TAG:
                    ((OrderConsumer) consumer).subscribe(topic, selectorExpression, new DefaultMessageListenerOrderly());
                    break;
//                case SQL92:
//                    ((OrderConsumer) consumer).subscribe(topic, MessageSelector.bySql(selectorExpression), new DefaultMessageListenerOrderly());
//                    break;
                default:
                    throw new IllegalArgumentException("Property 'selectorType' was wrong.");
            }
        } else {
            throw new IllegalArgumentException("Property 'consumeMode' was wrong.");
        }
    }

    private Properties getProperties() {
        Properties properties = new Properties();
        properties.setProperty(PropertyKeyConst.GROUP_ID, this.consumerGroup);
        properties.setProperty(PropertyKeyConst.AccessKey, this.applicationContext.getEnvironment().
                resolveRequiredPlaceholders(this.accessKey));
        properties.setProperty(PropertyKeyConst.SecretKey, this.applicationContext.getEnvironment().
                resolveRequiredPlaceholders(this.secretKey));
        properties.setProperty(PropertyKeyConst.ConsumeThreadNums, String.valueOf(this.consumeThreadMax));
        properties.put(PropertyKeyConst.NAMESRV_ADDR, this.nameServer);
        properties.setProperty(PropertyKeyConst.MessageModel, this.messageModel.getModeCN());
        properties.setProperty(PropertyKeyConst.ConsumeTimeout, String.valueOf(this.consumeTimeout));
        properties.setProperty(PropertyKeyConst.MaxReconsumeTimes, String.valueOf(this.maxReconsumeTimes));

        return properties;
    }

}
